package com.rt.gptnext.manager;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.ObjectUtil;
import com.rt.gptnext.common.constant.Constants;
import com.rt.gptnext.data.dto.CompletionDto;
import com.rt.gptnext.data.entity.OpenaiChat;
import com.rt.gptnext.service.OpenaiChatService;
import com.rt.gptnext.service.OpenaiCompletionService;
import jakarta.servlet.http.HttpServletResponse;
import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import org.hibernate.service.spi.ServiceException;
import org.springframework.stereotype.Service;

/**
 * @author gaoyiyang
 */
@Service
@AllArgsConstructor
public class OpenaiManager {

    private final OpenaiChatService openaiChatService;
    private final OpenaiCompletionService openaiCompletionService;

    private final static String HEADER_CHAT_ID = "chat-id";

    private final static String HEADER_QUESTION_ID = "question-id";

    private final static String HEADER_ANSWER_ID = "answer-id";

    /**
     * 开始对话
     *
     * @param completionDto 对话dto
     * @param response      servlet响应
     */
    public void completion(CompletionDto completionDto, HttpServletResponse response) {
        OpenaiChat chat = null;
        if (completionDto.getChatId() != null) {
            chat = openaiChatService.getById(completionDto.getChatId());
        }
        if (chat == null) {
            // 初始化聊天
            chat = openaiChatService.initChat(completionDto.getSystemPrompt());
        }
        // 构造问题
        var content = completionDto.getContent();
        executeCompletion(chat, response, completionDto.getStream(), content);
    }

    /**
     * 重新生成最后一个问题的对话
     *
     * @param completionDto 对话dto
     * @param response      servlet响应
     */
    public void resetCompletion(CompletionDto completionDto, HttpServletResponse response) {
        // 重置聊天
        var chat = openaiChatService.reset(completionDto.getChatId());
        var content = openaiCompletionService.reset(chat, chat.getVersion() - 1);
        executeCompletion(chat, response, completionDto.getStream(), ObjectUtil.defaultIfNull(content, Constants.EMPTY));
    }

    /**
     * 继续对话
     *
     * @param completionDto 对话dto
     * @param response      servlet响应
     */
    public void continueCompletion(CompletionDto completionDto, HttpServletResponse response) {
        // 获取聊天
        var chat = openaiChatService.getById(completionDto.getChatId());
        if (chat == null) {
            throw new ServiceException("聊天不存在");
        }
        var content = completionDto.getContent();
        executeCompletion(chat, response, completionDto.getStream(), content);
    }

    /**
     * 执行对话
     *
     * @param chat     聊天对象
     * @param response 响应
     * @param stream   是否流读取
     * @param content  问题
     */
    @SneakyThrows
    private void executeCompletion(OpenaiChat chat, HttpServletResponse response, Boolean stream, String content) {
        var questionId = IdUtil.fastSimpleUUID();
        var answerId = IdUtil.fastSimpleUUID();
        // 设置响应头
        response.addHeader(HEADER_CHAT_ID, chat.getId().toString());
        response.addHeader(HEADER_QUESTION_ID, questionId);
        response.addHeader(HEADER_ANSWER_ID, answerId);

        if (Boolean.FALSE.equals(stream)) {
            var answer = openaiCompletionService.completion(chat, content, questionId, answerId);
            response.setContentType("text/plain");
            IoUtil.writeUtf8(response.getOutputStream(), true, answer);
        } else {
            response.setContentType("text/event-stream");
            openaiCompletionService.streamCompletion(chat, content, response.getOutputStream(), questionId, answerId);
        }
    }

}
